#!/usr/bin/env python
# -*- coding: utf-8 -*-

import envs
import genpy
import rospy
import numpy as np

from std_msgs.msg import Float64
from control_msgs.msg import JointTrajectoryControllerState
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint

# 偏差阈值
BIAS_THRESHOLD = 0.05


class ArmTheta():
    def __init__(self):
        # 物体斜率订阅者
        self.object_theta_sub = rospy.Subscriber('/camera/image_raw/object_theta', Float64, self.object_theta_callback, queue_size=1)

        # 机械臂状态发布者
        self.joint_cmd_pub = rospy.Publisher('/xarm_controller/command', JointTrajectory, queue_size=10)

        # 机械臂状态订阅者
        self.joint_state_sub = rospy.Subscriber('/xarm_controller/state', JointTrajectoryControllerState, self.joint_state_callback, queue_size=10)

        # 机械臂关节初始化
        self.joint_states = []

        # 机械臂关节控制初始化
        self.joint_cmd_data = JointTrajectory()
        self.joint_cmd_data.joint_names = ['arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4', 'arm_joint_5']
        self.joint_cmd_data_point = JointTrajectoryPoint()
        self.joint_cmd_data_point.time_from_start = genpy.Duration(0.5)

        self.callback_times = 0
        self.callback_times_threshold = 1

        rospy.spin()

    def joint_state_callback(self, req):
        """
        机械臂状态回调函数
        """
        rospy.logdebug('joint state callback...')
        rospy.logdebug('req.actual.positions: %s', req.actual.positions)

        self.joint_states = list(req.actual.positions)

        rospy.logdebug('joint state callback OK')

    def object_theta_callback(self, data):
        self.callback_times += 1
        if self.callback_times == self.callback_times_threshold:
            self.callback_times = 0
            rospy.logdebug('object theta callback...')
            object_theta = data.data
            rospy.logdebug('object theta: %s', object_theta)

            # 计算偏差
            if object_theta < np.pi / 2:
                bias_x = round(object_theta - self.joint_states[4], 3)
            else:
                bias_x = round(object_theta - np.pi - self.joint_states[4], 3)

            if bias_x < BIAS_THRESHOLD and bias_x > (-BIAS_THRESHOLD):
                bias_x = 0

            # 比例控制信号计算
            control_x = 0.8 * bias_x
            rospy.logdebug('control x: %s', control_x)

            # 调整关节状态
            self.joint_cmd_data_point.positions = self.joint_states
            self.joint_cmd_data_point.positions[4] = self.joint_states[4] + control_x
            rospy.logdebug('publish data is %s', self.joint_cmd_data_point.positions)
            self.joint_cmd_data.points = [self.joint_cmd_data_point]

            # 发送控制信号
            self.joint_cmd_pub.publish(self.joint_cmd_data)
            rospy.logdebug('publish data OK')

            rospy.logdebug('object theta callback OK')


if __name__ == '__main__':
    rospy.init_node('npu_robot_arm_theta', anonymous=True, log_level=envs.log_level)
    arm_theta = ArmTheta()
